# Bezier curve

import matplotlib.pyplot as plt
import numpy as np
import scipy.special

show_animation = True

# 对于n_points个离散点，等间隔t
def calc_bezier_path(control_points, n_points):
    traj = []
    for t in np.linspace(0, 1, n_points):
        traj.append(bezier(t, control_points))
    return np.array(traj)

# 贝塞尔曲线的每一项
def bernstein_poly(n, i, t):
    return scipy.special.comb(n, i) * t ** i * (1 - t) ** (n - i)

# 贝塞尔曲线
def bezier(t, control_points):
    n = len(control_points) - 1
    # B(t) = P_0*(1-t)^3 + 3*P_1*t*(1-t)^2 + 3*P_2*t^2*(1-t) + P_3*t^3  t∈[0,1]
    return np.sum([bernstein_poly(n, i, t) * control_points[i] for i in range(n + 1)], axis=0)

# 计算贝塞尔曲线在控制点的导数
def bezier_derivatives_control_points(control_points, n_derivatives):
    w = {0: control_points}
    for i in range(n_derivatives):
        n = len(w[i])
        w[i + 1] = np.array([(n - 1) * (w[i][j + 1] - w[i][j])
                             for j in range(n - 1)])
    return w

# 计算曲率
def curvature(dx, dy, ddx, ddy):
    return (dx * ddy - dy * ddx) / (dx ** 2 + dy ** 2) ** (3 / 2)

def main():
    # 三阶贝塞尔曲线由4个点确定
    control_points = np.array([[5., 1.], [-2.78, 1.], [-11.5, -4.5], [-6., -8.]])
    path = calc_bezier_path(control_points, n_points=100)
    # Display the tangent, normal and radius of cruvature at a given point
    t = 0.5  # Number in [0, 1]
    x_target, y_target = bezier(t, control_points)
    derivatives_cp = bezier_derivatives_control_points(control_points, 2)
    point = bezier(t, control_points)
    dt = bezier(t, derivatives_cp[1])
    ddt = bezier(t, derivatives_cp[2])
    # Radius of curvature
    radius = 1 / curvature(dt[0], dt[1], ddt[0], ddt[1])
    # Normalize derivative
    dt /= np.linalg.norm(dt, 2)
    tangent = np.array([point, point + dt])
    normal = np.array([point, point + [- dt[1], dt[0]]])
    curvature_center = point + np.array([- dt[1], dt[0]]) * radius
    circle = plt.Circle(tuple(curvature_center), radius,
                        color=(0, 0.8, 0.8), fill=False, linewidth=1)
    if show_animation:  # pragma: no cover
        fig, ax = plt.subplots()
        ax.plot(path.T[0], path.T[1], label="Bezier Path")
        ax.plot(control_points.T[0], control_points.T[1],
                '--o', label="Control Points")
        ax.plot(x_target, y_target)
        ax.plot(tangent[:, 0], tangent[:, 1], label="Tangent")
        ax.plot(normal[:, 0], normal[:, 1], label="Normal")
        ax.add_artist(circle)
        ax.legend(loc='best')
        ax.axis("equal")
        ax.grid(True)
        plt.show()

if __name__ == '__main__':
    main()